#include "Skeleton.h"
#include "Model.h"
#include "File.h"

Skeleton::Skeleton(std::string name, std::shared_ptr<SkeletonAsset> asset, SkeletonManager* manager)
{
    this->name = name;
    this->skeletonAsset = asset;
    this->skeletonManager = manager;
}

void Skeleton::UpdateInner(int id, std::shared_ptr<std::vector<glm::mat4>> anim)
{
    if (arrayBonesFlag[id] == true)
        return;

    if (arrayBones[id].parent == -1)
    {
        if ((*anim).size() > arrayBones[id].animId && arrayBones[id].animId != -1)
            arrayBones[id].UpdateOffset(offset, (*anim)[arrayBones[id].animId]);
        else
            arrayBones[id].UpdateOffset(offset);
        arrayBonesFlag[id] = true;
        CurrentBindPose[id] = arrayBones[id].offset * arrayBones[id].bindOffset;
        CurrentPose[id] = arrayBones[id].offset;
    }
    else{
        UpdateInner(arrayBones[id].parent, anim);
        if ((*anim).size() > arrayBones[id].animId && arrayBones[id].animId != -1)
            arrayBones[id].UpdateOffset(arrayBones[arrayBones[id].parent].offset, (*anim)[arrayBones[id].animId]);
        else
            arrayBones[id].UpdateOffset(arrayBones[arrayBones[id].parent].offset);
        arrayBonesFlag[id] = true;
        CurrentBindPose[id] = arrayBones[id].offset * arrayBones[id].bindOffset;
        CurrentPose[id] = arrayBones[id].offset;
    }
}

void Skeleton::UpdateOffset(std::shared_ptr<std::vector<glm::mat4>> anim)
{
    for (int i = 0; i < arrayBonesFlag.size(); i++)
    {
        arrayBonesFlag[i] = false;
    }
    // memset((arrayBonesFlag.data()), 0, sizeof(bool) * arrayBonesFlag.size());

    // rootBone->UpdateOffset(anim, CurrentBindPose, CurrentPose, offset);
    for (int i=0;i<arrayBones.size();i++){
        UpdateInner(i, anim);
    }
}

static glm::mat4 MixSRT(const SRT &A, const SRT &B, float blend)
{
    if (blend < 0.01f)
    {
        glm::mat4 m = glm::scale(glm::mat4_cast(A.rotation), A.scale);
        m[3] = glm::vec4(A.translation, 1.0f);
        return m;
    }
    else if (blend > 0.99f)
    {
        glm::mat4 m = glm::scale(glm::mat4_cast(B.rotation), B.scale);
        m[3] = glm::vec4(B.translation, 1.0f);
        return m;
    }

    glm::vec3 s = glm::mix(A.scale, B.scale, blend);
    glm::vec3 t = glm::mix(A.translation, B.translation, blend);
    glm::quat q = glm::slerp(A.rotation, B.rotation, blend);
    // return SRT(s,q,t);
    glm::mat4 m = glm::scale(glm::mat4_cast(q), s);
    m[3] = glm::vec4(t, 1.0f);
    return m;
}

void SRTFrametoMat4Frame(std::shared_ptr<std::vector<SRT>> frame, std::shared_ptr<std::vector<glm::mat4>> out)
{
    for (int i = 0; i < out->size(); i++)
    {
        (*out)[i] = SRT::SRTtoMat4((*frame)[i]);
    }
}

void BlendFrameSRT(std::shared_ptr<std::vector<SRT>> frameA,
                   std::shared_ptr<std::vector<SRT>> frameB,
                   float alpha,
                   std::shared_ptr<std::vector<glm::mat4>> out)
{
    for (uint32_t i = 0; i < out->size(); i++)
    {
        (*out)[i] = MixSRT((*frameA)[i], (*frameB)[i], alpha);
    }
}

void Skeleton::Play2BlendAnimation(const std::string &animName, const std::string &animName2, uint32_t tick, float blend)
{
    if (skeletonAsset->orderedAnimations.find(animName) == skeletonAsset->orderedAnimations.end())
    {
        throw std::runtime_error("Animation " + animName + " not found");
    }
    if (skeletonAsset->orderedAnimations.find(animName2) == skeletonAsset->orderedAnimations.end())
    {
        throw std::runtime_error("Animation " + animName2 + " not found");
    }

    auto anim = skeletonAsset->orderedAnimations[animName];
    auto anim2 = skeletonAsset->orderedAnimations[animName2];

    auto frameId = tick % anim->GetFrameSize();
    auto frameId2 = tick % anim2->GetFrameSize();

    // 这里只取需要播放动画的骨骼的数量
    auto boneTransforms = std::make_shared<std::vector<glm::mat4>>(skeletonAsset->animBoneNameIndex.size());

    BlendFrameSRT(anim->GetFrameAnim(frameId),
                  anim2->GetFrameAnim(frameId2),
                  blend, boneTransforms);

    UpdateOffset(boneTransforms);
}

void Skeleton::PlayAnimation(const std::string &animName, uint32_t tick)
{
    if (skeletonAsset->orderedAnimations.find(animName) == skeletonAsset->orderedAnimations.end())
    {
        throw std::runtime_error("Animation " + animName + " not found");
    }

    auto anim = skeletonAsset->orderedAnimations[animName];

    auto frameId = tick % anim->GetFrameSize();

    // 这里只取需要播放动画的骨骼的数量, 把anim中超过数量的骨骼动画剔除
    // anim->GetFrameAnim(frameId)
    auto boneTransforms = std::make_shared<std::vector<glm::mat4>>(skeletonAsset->animBoneNameIndex.size());
    SRTFrametoMat4Frame(anim->GetFrameAnim(frameId), boneTransforms);

    UpdateOffset(boneTransforms);
}

void Skeleton::PrepareToDraw(){
    if (!toDraw)
        return;
    drawId = skeletonManager->instanceDrawCount;
    skeletonManager->instanceDrawCount++;
}

void Skeleton::ProcessManagerData()
{
    if (!toDraw)
        return;
    if (skinBonesIndices.size() == 0)
        return;
    int dataWidth = skeletonManager->width * 4;
    int i = 0;
    int y = drawId * dataWidth;

    if (skeletonManager->useDQBSkinning)
    {
        glm::quat r;
        glm::vec3 t, s;

        glm::fdualquat dq;

        for (int x = 0; x < dataWidth; x += 12)
        {
            s = ExtractScale(CurrentBindPose[skinBonesIndices[i]]);
            r = glm::quat_cast(ExtractRotationMatrix(CurrentBindPose[skinBonesIndices[i]], s));
            t = glm::vec3(CurrentBindPose[skinBonesIndices[i]][3]);
            dq[0] = r;
            dq[1] = glm::quat(0, t.x, t.y, t.z) * r * .5f;

            skeletonManager->AnimTransformData[y + x + 0] = dq[0][0];
            skeletonManager->AnimTransformData[y + x + 1] = dq[0][1];
            skeletonManager->AnimTransformData[y + x + 2] = dq[0][2];
            skeletonManager->AnimTransformData[y + x + 3] = dq[0][3];
            skeletonManager->AnimTransformData[y + x + 4] = dq[1][0];
            skeletonManager->AnimTransformData[y + x + 5] = dq[1][1];
            skeletonManager->AnimTransformData[y + x + 6] = dq[1][2];
            skeletonManager->AnimTransformData[y + x + 7] = dq[1][3];
            skeletonManager->AnimTransformData[y + x + 8] = s[0];
            skeletonManager->AnimTransformData[y + x + 9] = s[1];
            skeletonManager->AnimTransformData[y + x + 10] = s[2];
            // skeletonManager->AnimTransformData[y + x + 11] = 1.0f;
            i++;
            if (i >= skinBonesIndices.size())
            {
                break;
            }
        }
    }
    else
    {
        for (int x = 0; x < dataWidth; x += 16)
        {
            auto m = CurrentBindPose[skinBonesIndices[i]];
            skeletonManager->AnimTransformData[y + x + 0] = m[0][0];
            skeletonManager->AnimTransformData[y + x + 1] = m[0][1];
            skeletonManager->AnimTransformData[y + x + 2] = m[0][2];
            skeletonManager->AnimTransformData[y + x + 3] = m[0][3];
            skeletonManager->AnimTransformData[y + x + 4] = m[1][0];
            skeletonManager->AnimTransformData[y + x + 5] = m[1][1];
            skeletonManager->AnimTransformData[y + x + 6] = m[1][2];
            skeletonManager->AnimTransformData[y + x + 7] = m[1][3];
            skeletonManager->AnimTransformData[y + x + 8] = m[2][0];
            skeletonManager->AnimTransformData[y + x + 9] = m[2][1];
            skeletonManager->AnimTransformData[y + x + 10] = m[2][2];
            skeletonManager->AnimTransformData[y + x + 11] = m[2][3];
            skeletonManager->AnimTransformData[y + x + 12] = m[3][0];
            skeletonManager->AnimTransformData[y + x + 13] = m[3][1];
            skeletonManager->AnimTransformData[y + x + 14] = m[3][2];
            skeletonManager->AnimTransformData[y + x + 15] = m[3][3];
            i++;
            if (i >= skinBonesIndices.size())
            {
                break;
            }
        }
    }
}


void Skeleton::ProcessAllMeshInstance()
{
    if (!toDraw)
        return;

    for (auto mesh : bindMeshes)
    {
        mesh->drawId = drawId;
        mesh->drawMask = BodyMask;
        mesh->ToDraw();
    }
}


SkeletonAsset::SkeletonAsset(std::string name, const aiNode *rootBone, ModelAsset *model)
{
    this->SkeletonType = name;
    this->rootBoneName = rootBone->mName.C_Str();

    CollectBone(rootBone, nullptr, model);
    
    // 遍历骨骼，为每个骨骼重新设置id
    // 优先设置animated的骨骼，把它们的id排序在前面
    std::list<std::string> animBones;
    std::list<std::string> skinBones;

    std::list<std::string> posBones;
    for (auto &bone : bones)
    {
        auto flag = false;
        if (bone.second->animated)
        {
            animBones.push_back(bone.second->name);
            flag =true;
        }
        
        if (bone.second->skinned)
        {
            skinBones.push_back(bone.second->name);
            flag = true;
        }

        if (!flag)
        {
            posBones.push_back(bone.second->name);
        }
    }

    // id重置，这是为了防止id有负值
    int id = 0;
    for (auto &bone : bones)
    {
        bone.second->id = id;
        boneNameIndex.insert(std::make_pair(bone.first, id));
        id++;
    }

    int animid = 0;
    int skinid = 0;
    // id从0开始，先更新animBones中的骨骼的id，再更新nonAnimBones中的骨骼的id
    for (auto &boneName : animBones)
    {
        bones[boneName]->animId = animid;
        animBoneNameIndex.insert(std::make_pair(boneName, animid));
        animid++;
    }

    for (auto &boneName : skinBones)
    {
        bones[boneName]->skinId = skinid;
        skinBoneNameIndex.insert(std::make_pair(boneName, skinid));
        skinid++;
    }
}

void SkeletonAsset::CollectBone(const aiNode *node, const aiNode *parent, ModelAsset *model)
{
    std::string name = node->mName.C_Str();
    int id = model->GetBoneIndex(name);
    glm::mat4 bindOffset = glm::mat4(1.0f);

    // id = -1代表它不是个正常骨骼，需要判断是否要保留
    if (id == -1)
    {
        // 如果它没有子节点，或者子节点都不是骨骼（id = -1），则不保留
        if (node->mNumChildren == 0)
        {
            return;
        }
        bool hasBone = false;
        // 如果它有子节点，则遍历子节点，看是否有骨骼
        for (int i = 0; i < node->mNumChildren; i++)
        {
            if (model->GetBoneIndex(node->mChildren[i]->mName.C_Str()) != -1)
            {
                hasBone = true;
                break;
            }
        }
        if (!hasBone)
        {
            return;
        }
    }
    else
        bindOffset = model->boneOffsetCache[id];

    glm::mat4 restOffset = aiMatrix4_cast(node->mTransformation);
    glm::mat4 baseRestOffsetInverse = glm::inverse(restOffset);
    std::string parentName;
    parentName.clear();

    // 我们不应该在最后collectBone的时候更新父子关系，首先得确定子节点是个骨骼，才能做这件事
    if (parent != nullptr)
    {
        parentName = parent->mName.C_Str();
        if (bones.find(parentName) == bones.end())
        {
            // 报错
            throw std::runtime_error("Skeleton::CollectBone: parent bone not found");
        }

        bones.find(parentName)->second->childrenNames.push_back(name);
    }

    std::shared_ptr<BoneAsset> bone = std::make_shared<BoneAsset>(name, id, parentName, restOffset, bindOffset, baseRestOffsetInverse);
    bones.insert(std::make_pair(name, bone));

    for (uint32_t i = 0; i < node->mNumChildren; i++)
        CollectBone(node->mChildren[i], node, model);
}


// var name = node.Key;
// var id = ReadYamlInfo.LoadField(info, "ID", -1);
// var skinid = ReadYamlInfo.LoadField(info, "SkinID", -1);
// var animId = ReadYamlInfo.LoadField(info, "AnimID", -1);
// var parentName = ReadYamlInfo.LoadField(info, "Parent", "NO_Parent");
// var parentID = ReadYamlInfo.LoadField(info, "ParentID", -1);
// mat4 restPose = ReadYamlInfo.LoadMat4(info, "RestPose");
// mat4 restPoseInv = ReadYamlInfo.LoadMat4(info, "RestPoseInv");
// mat4 bindPose = ReadYamlInfo.LoadMat4(info, "BindPose");

void SkeletonAsset::Export(std::string filename)
{
    YamlWriter writer(filename);

    for (auto &kv : bones){
        auto bone = kv.second;
        writer.WriteLine(bone->name + ":");
        writer.WriteLineAfterTab("ID: " + std::to_string(bone->id));
        if (bone->skinned)
            writer.WriteLineAfterTab("Skin: true");
        if (bone->animated)
            writer.WriteLineAfterTab("Anim: true");
        if (bone->parentName.empty()){
            writer.WriteLineAfterTab("Parent: NO_Parent");
        }
        else{
            writer.WriteLineAfterTab("Parent: " + bone->parentName);
            writer.WriteLineAfterTab("ParentID: " + std::to_string(boneNameIndex[bone->parentName]));
        }
        writer.WriteMat4("RestPose", bone->restPose);
        writer.WriteMat4("RestPoseInv", bone->restPoseInv);
        writer.WriteMat4("BindPose", bone->bindPose);
    }
}

SkeletonManager::SkeletonManager(std::shared_ptr<SkeletonAsset> asset)
{
    this->asset = asset;
    AnimTransformData.resize(width * height * 4);
    glPixelStorei(GL_UNPACK_ALIGNMENT, 1);
    glGenTextures(1, &AnimTransformTexture);
    glBindTexture(GL_TEXTURE_2D, AnimTransformTexture);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA32F, width, height, 0, GL_RGBA, GL_FLOAT, NULL);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
}

void SkeletonManager::UpdateAnimTexture()
{
    instanceDrawCount = 0;
    glPixelStorei(GL_UNPACK_ALIGNMENT, 1);
    glBindTexture(GL_TEXTURE_2D, AnimTransformTexture);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA32F, width, height, 0, GL_RGBA, GL_FLOAT, AnimTransformData.data());
}

void SkeletonManager::PrepareInstanceDraw(const gl::IShader &shader) const
{
    shader.setInt("boneAnimTexture", 5);
    glActiveTexture(GL_TEXTURE5);
    glBindTexture(GL_TEXTURE_2D, AnimTransformTexture);
    glActiveTexture(GL_TEXTURE0);
    shader.setBool("useDQB", useDQBSkinning);
}

std::shared_ptr<Skeleton> SkeletonManager::CreateSkeleton()
{
    std::shared_ptr<Skeleton> sk = std::make_shared<Skeleton>(asset->SkeletonType, asset, this);
    sk->arrayBones.resize(asset->bones.size());
    sk->arrayBonesFlag.resize(asset->bones.size());
    sk->animBonesIndices.resize(asset->animBoneNameIndex.size());
    sk->skinBonesIndices.resize(asset->skinBoneNameIndex.size());

    // 遍历bones创建骨骼
    for (auto &bone : asset->bones)
    {
        if (bone.second->parentName.empty())
        {
            sk->arrayBones[bone.second->id] = ArrayBone(bone.second->animId, -1, bone.second->restPose, bone.second->bindPose, bone.second->restPoseInv);
        }
        else
        {
            sk->arrayBones[bone.second->id] = ArrayBone(bone.second->animId, asset->bones[bone.second->parentName]->id, bone.second->restPose, bone.second->bindPose, bone.second->restPoseInv);
        }
    }

    for (auto &bone : asset->skinBoneNameIndex)
    {
        sk->skinBonesIndices[bone.second] = asset->bones[bone.first]->id;
    }

    for (auto &bone : asset->animBoneNameIndex)
    {
        sk->animBonesIndices[bone.second] = asset->bones[bone.first]->id;
    }

    // sk->allBones.resize(asset->bones.size());

    // // 遍历bones创建骨骼
    // for (auto &bone : asset->bones)
    // {
    //     sk->allBones[bone.second->id] = std::make_shared<Bone>(bone.second);
    // }
    // // 遍历bones创建骨骼的父子关系
    // for (auto &bone : asset->bones)
    // {
    //     if (!bone.second->parentName.empty())
    //     {
    //         sk->allBones[bone.second->id]->parent = sk->allBones[asset->bones[bone.second->parentName]->id];
    //         sk->allBones[asset->bones[bone.second->parentName]->id]->children.push_back(sk->allBones[bone.second->id]);
    //     }
    // }
    // sk->rootBone = sk->allBones[asset->bones[asset->rootBoneName]->id];

    sk->CurrentPose.resize(asset->bones.size());
    sk->CurrentBindPose.resize(asset->bones.size());

    skeletons.push_back(sk);

    return sk;
}